from gurobipy import *
import numpy as np
from traffic_engineer import *
from robust_multi_cluster_scaleup_weighted_traffic_engineer import *
from traffic_engineering_class import *

# same implementation as RobustMultiClusterTrafficEngineer but we weigh traffic matrix with its max scale up factor
class RobustMultiClusterScaleUpWeightedTrafficEngineerHeuristic(RobustMultiClusterScaleUpWeightedTrafficEngineer):
	## Takes in two new variables:
	## update_frequency - how much time in-between every traffic engineering events
	def __init__(self, aurora_network, all_interblock_paths, update_frequency, training_length, num_k, reduce_multihop=False, ignore_flows=None):
		RobustMultiClusterScaleUpWeightedTrafficEngineer.__init__(self, aurora_network, all_interblock_paths, update_frequency, training_length, num_k, reduce_multihop=reduce_multihop)
		self.flow_pairs_to_ignore = ignore_flows
		self.update_frequency = update_frequency
		self.training_length = training_length
		self.num_k = num_k
		return

	def get_filename_param(self):
		return "robustweightedheuristicte_r{}t{}c{}".format(self.update_frequency, self.training_length, self.num_k)

	def get_num_clusters_for_training(self):
		return self.num_k
	# returns a string description of this class
	def get_string(self):
		return "Robust scale up weighted heuristic TE, {} clusters, w/ update_frequency: {} and training_length: {}".format(self.num_k, self.update_frequency, self.training_length)

	def traffic_engineering_class(self):
		return TrafficEngineeringClass.ROBUST_WEIGHTED_TE

	## Currently does not involve another step that reduces hop count
	def _compute_max_traffic_scaleup(self, logical_topology_adj_matrix, traffic_matrix):
		nblocks = self.aurora_network.get_num_blocks()
		model = Model("Find the max traffic scale up")
		model.setParam( 'OutputFlag', False )
		mu = model.addVar(lb=0., ub=GRB.INFINITY, obj=0, vtype=GRB.CONTINUOUS, name="mu")
		routing_weight_vars = {}
		link_capacity_constraints = [None] * nblocks
		for i in range(nblocks):
			link_capacity_constraints[i] = [None] * nblocks
			for j in range(nblocks):
				if i != j:
					link_capacity_constraints[i][j] = LinExpr()
		for i in range(nblocks):
			for j in range(nblocks):
				if i != j:
					weight_constraint = LinExpr()
					for path in self.all_interblock_paths[i][j]:
						routing_weight_vars[path] = model.addVar(lb=0., ub=GRB.INFINITY, obj=1., vtype=GRB.CONTINUOUS, name="w_{}".format(path))
						weight_constraint += routing_weight_vars[path]
					model.addConstr(lhs=mu * traffic_matrix[i][j], sense=GRB.LESS_EQUAL, rhs=weight_constraint)

		# stage 1: setup all the optimization variables that are the routing weights
		# stage 1.5 : also adds the traffic flow satisfiability constraints
		for i in range(nblocks):
			for j in range(nblocks):
				if i != j:
					for path in self.all_interblock_paths[i][j]:
						src = path[0]
						for path_hop in range(1, len(path), 1):
							dst = path[path_hop]
							link_capacity_constraints[src][dst] += (routing_weight_vars[path])
							src = dst
		## stage 2 : add the link utilization constraints to the model
		for i in range(nblocks):
			for j in range(nblocks):
				if i != j:
					lu_constraint = link_capacity_constraints[i][j]
					capacity = min(self.aurora_network.get_link_capacity(i), self.aurora_network.get_link_capacity(j))
					model.addConstr(lhs=lu_constraint, sense=GRB.LESS_EQUAL, rhs=capacity * logical_topology_adj_matrix[i][j])
		# stage 3: set the objective
		model.setObjective(mu, GRB.MAXIMIZE)
		try: 
			model.optimize()
			return mu.x
		except GurobiError as e:
			print ("Error code " + str(e. errno ) + ": " + str(e))
			return None
		except AttributeError :
			print ("Encountered an attribute error ")
			return None


	def compute_path_weights(self, logical_topology_adj_matrix, traffic_matrices, spare_traffic=0.1, fraction_threshold=0.001, cutoff=0.002):
		## newly added for robustness
		nblocks = self.aurora_network.get_num_blocks()
		for traffic_matrix in traffic_matrices:
			traffic_sum = sum([sum(x) for x in traffic_matrix])
			reset_num_entries = 0
			for i in range(nblocks):
				for j in range(nblocks):
					if i != j and self.flow_pairs_to_ignore is not None and (i,j) not in self.flow_pairs_to_ignore:
						if traffic_matrix[i][j] < fraction_threshold * traffic_sum:
							reset_num_entries += 1
			if reset_num_entries > 0:
				scratch_traffic = (spare_traffic * traffic_sum) / (1 - spare_traffic) / float(reset_num_entries)
				for i in range(nblocks):
					for j in range(nblocks):
						if i != j and self.flow_pairs_to_ignore is not None and (i,j) not in self.flow_pairs_to_ignore:
							if traffic_matrix[i][j] < fraction_threshold * traffic_sum:
								traffic_matrix[i][j] = scratch_traffic
		#for traffic_matrix in traffic_matrices:
		#	traffic_sum = sum([sum(x) for x in traffic_matrix])
		#	reset_num_entries = 0
		#	for i in range(nblocks):
		#		for j in range(nblocks):
		#			if i != j:
		#				if traffic_matrix[i][j] < 0.0001:
		#					reset_num_entries += 1
		#	for i in range(nblocks):
		#		for j in range(nblocks):
		#			if i != j:
		#				if traffic_matrix[i][j] < 0.0001:
		#					traffic_matrix[i][j] = (spare_traffic * traffic_sum) / (1 - spare_traffic) / float(reset_num_entries)
		## newly added for robustness
		for tm_index in range(len(traffic_matrices)):
			scaleup_factor = self._compute_max_traffic_scaleup(logical_topology_adj_matrix, traffic_matrices[tm_index])
			# now, scale each entry in traffic matrix by scaleup factor
			for i in range(nblocks):
				for j in range(nblocks):
					if i != j:
						traffic_matrices[tm_index][i][j] *= scaleup_factor
		routing_weights, worst_case_mlu = self._compute_path_weights_minimize_worstcase_mlu(logical_topology_adj_matrix, traffic_matrices)
		if self.reduce_multihop:
			routing_weights = self._minimize_multihop(logical_topology_adj_matrix, traffic_matrices, worst_case_mlu)
		#### Jason Newlu added hack (BEGIN)
		for i in range(nblocks):
			for j in range(nblocks):
				if i != j:
					max_entry = 0.
					for k in range(len(traffic_matrices)):
						max_entry = max(max_entry, traffic_matrices[k][i][j])
					if max_entry <= cutoff:
						path_total_capacity = 0.
						for path in self.all_interblock_paths[i][j]:
							curr_node = path[0]
							path_capacity = 10E15
							for index in range(1, len(path), 1):
								next_node = path[index]
								cap = min(self.aurora_network.get_link_capacity(curr_node), self.aurora_network.get_link_capacity(next_node))
								path_capacity = min(path_capacity, logical_topology_adj_matrix[curr_node][next_node] * cap)
								curr_node = next_node
							path_total_capacity += path_capacity
							routing_weights[path] = path_capacity
						for path in self.all_interblock_paths[i][j]:
							routing_weights[path] /= path_total_capacity
		#### Jason Newlu added hack (END)
		return routing_weights